import torch

# root用户从本地加载
source = '/root/.cache/torch/hub/huggingface_pytorch-transformers_main'
# 直接使用预训练的bert中文模型
model_name = 'bert-base-chinese'
# 通过torch.hub获得已经训练好的bert-base-chinese模型
model = torch.hub.load(source, 'model', model_name, source='local')
# 获得对应的字符映射器，它将把中文的每个字映射成一个数字
tokenizer = torch.hub.load(source, 'tokenizer', model_name, source='local')

def get_bert_endcode_for_single(text:str):
    """
    使用bert base Chinese对接收的字符串进行编码，是针对每个字进行编码
    :param text: 接收的字符串
    :return: 编码后的向量
    """
    indexed_token = tokenizer.encode(text)[1:-1]
    # 每句话输出的编码中会包含开头和结尾的标识符，101 102 对多段话有意义，对一句话没什么用
    # print(indexed_token)

    # 把上一步的列表转成张量
    token_tensor = torch.LongTensor([indexed_token])

    with torch.no_grad():
        encode_layers = model(token_tensor)

    return encode_layers[0] # 这个结果是几维度？，现在是3维，这个下标0的意思就是取last_hidden_state结果

if __name__ == '__main__':
    print(get_bert_endcode_for_single("今天好热啊"))